/*
 * Copyright (C) 2022  Rockchip Electronics Co., Ltd.
 * Authors:
 *     YuQiaowei <cerf.yu@rock-chips.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <iostream>
#include <fstream>
#include <sstream>
#include <cstddef>
#include <cmath>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <unistd.h>
#include <pthread.h>

#include "im2d.h"
#include "RgaUtils.h"
#include "rga.h"         // rga_info_t, RK_FORMAT_*
#include "RgaApi.h"     // c_RkRgaBlit, c_RkRgaInit, c_RkRgaDeInit
#include "dma_alloc.h"   // dma_buf_alloc, dma_buf_free

// 对齐工具（RGA 通常要求 stride 按 16 对齐）
static inline int align_up(int x, int a) { return ((x + a - 1) / a) * a; }

// 全局静态输出 buffer（复用）
static void* g_dst_virt = nullptr;
static int   g_dst_fd   = -1;
static size_t g_dst_size = 0;
static int   g_max_w = 0;
static int   g_max_h = 0;

// 线程锁保护 RGA 调用
static pthread_mutex_t g_rga_lock = PTHREAD_MUTEX_INITIALIZER;
static bool g_rga_inited = false;

// 初始化：设置最大分辨率并分配 dst dma buffer（RGB888）
// max_width, max_height: 需要支持的最大分辨率（例如 3840,2160）
extern "C" int nv12_to_rgb_init(int max_width, int max_height) {
    if (max_width <= 0 || max_height <= 0) return -1;

    if (g_rga_inited) {
        // 如果已初始化，但新请求的大小更大则尝试重新分配
        if (max_width <= g_max_w && max_height <= g_max_h) return 0;
        // 否则先释放再重新分配
        if (g_dst_virt) {
            dma_buf_free((int)g_dst_size, &g_dst_fd, g_dst_virt);
            g_dst_virt = nullptr;
            g_dst_fd = -1;
            g_dst_size = 0;
        }
        g_rga_inited = false;
    }

    // 计算需要的 dst buffer 大小（使用 stride 对齐）
    int wstride = align_up(max_width, 16);
    int hstride = max_height;
    size_t dst_size = (size_t)wstride * (size_t)hstride * 3; // RGB888

    // 分配 dma32_heap（小于4G）
    int ret = dma_buf_alloc(DMA_HEAP_DMA32_UNCACHED_PATH, dst_size, &g_dst_fd, &g_dst_virt);
    if (ret < 0 || g_dst_virt == nullptr) {
        fprintf(stderr, "nv12_to_rgb_init: dma_buf_alloc(dst) failed ret=%d\n", ret);
        return -1;
    }
    // 清零
    memset(g_dst_virt, 0, dst_size);

    // init RGA (idempotent)
    if (c_RkRgaInit() < 0) {
        fprintf(stderr, "nv12_to_rgb_init: c_RkRgaInit failed\n");
        dma_buf_free((int)dst_size, &g_dst_fd, g_dst_virt);
        g_dst_virt = nullptr;
        g_dst_fd = -1;
        g_dst_size = 0;
        return -1;
    }

    g_dst_size = dst_size;
    g_max_w = max_width;
    g_max_h = max_height;
    g_rga_inited = true;

    return 0;
}

// 为 Python/调用方分配一个 src NV12 dma_buf（返回 fd 和 虚拟地址）
extern "C" int nv12_alloc_src(int width, int height, int* out_fd, void** out_virt) {
    if (!out_fd || !out_virt) return -1;
    if (width <= 0 || height <= 0) return -1;

    // 计算 NV12 所需大小（按 stride 对齐）
    int wstride = align_up(width, 16);
    int hstride = height;
    size_t src_size = (size_t)wstride * (size_t)hstride + (size_t)wstride * (size_t)hstride / 2; // NV12: 1.5
    int fd = -1;
    void* virt = nullptr;
    int ret = dma_buf_alloc(DMA_HEAP_DMA32_UNCACHED_PATH, src_size, &fd, &virt);
    if (ret < 0 || virt == nullptr) {
        fprintf(stderr, "nv12_alloc_src: dma_buf_alloc(src) failed ret=%d\n", ret);
        return -1;
    }
    // zero
    memset(virt, 0, src_size);

    *out_fd = fd;
    *out_virt = virt;
    return (int)src_size; // return allocated size (positive) for convenience
}

// 释放 src buffer（与 nv12_alloc_src 配对）
extern "C" void nv12_free_src(int src_fd, size_t src_size, void* src_virt) {
    if (src_virt) {
        dma_buf_free((int)src_size, &src_fd, src_virt);
    }
}

// 主转换函数：给定 src_fd、width、height，把 NV12 转成 RGB888（输出写入全局 dst）
// stride 自动按 16 对齐。线程安全（blit 处加锁）。
// 注意：src 必须是 NV12（YCbCr_420_SP），并且调用方保证 src buffer 实际已填入 NV12 数据。
// 返回 0 成功，否则负值。
extern "C" int nv12_to_rgb_convert(int src_fd, int width, int height) {
    if (!g_rga_inited) return -1;
    if (width <= 0 || height <= 0) return -1;
    if (width > g_max_w || height > g_max_h) {
        fprintf(stderr, "nv12_to_rgb_convert: requested size (%d,%d) > initialized max (%d,%d)\n",
                width, height, g_max_w, g_max_h);
        return -1;
    }

    // 设置 src/dst rga_info_t
    rga_info_t src_info;
    rga_info_t dst_info;
    memset(&src_info, 0, sizeof(src_info));
    memset(&dst_info, 0, sizeof(dst_info));

    // stride 要按芯片和实现对齐（通常16）
    int wstride_src = align_up(width, 16);
    int hstride_src = height;

    int wstride_dst = align_up(width, 16);
    int hstride_dst = height;

    // 填写 src_info
    src_info.fd = src_fd;
    src_info.mmuFlag = 1; // 使用 IOMMU/缓存? (取决平台). 保留1让驱动处理
    // rect 信息
    src_info.rect.xoffset = 0;
    src_info.rect.yoffset = 0;
    src_info.rect.width = width;
    src_info.rect.height = height;
    src_info.rect.wstride = wstride_src;
    src_info.rect.hstride = hstride_src;
    src_info.rect.format = RK_FORMAT_YCbCr_420_SP; // NV12

    // 填写 dst_info
    dst_info.fd = g_dst_fd;
    dst_info.mmuFlag = 1;
    dst_info.rect.xoffset = 0;
    dst_info.rect.yoffset = 0;
    dst_info.rect.width = width;
    dst_info.rect.height = height;
    dst_info.rect.wstride = wstride_dst;
    dst_info.rect.hstride = hstride_dst;
    dst_info.rect.format = RK_FORMAT_RGB_888; // RGB888

    // 如果需要 color-space 或其它转换选项，可在 info 的其它字段设置（此处使用驱动默认）
    int ret = 0;
    pthread_mutex_lock(&g_rga_lock);
    ret = c_RkRgaBlit(&src_info, &dst_info, NULL);
    pthread_mutex_unlock(&g_rga_lock);

    if (ret != 0) {
        fprintf(stderr, "nv12_to_rgb_convert: c_RkRgaBlit failed ret=%d\n", ret);
        return -1;
    }

    return 0;
}

// 获取 dst 的 dma fd（供 Python 传给其他进程或 mmap 用）
extern "C" int nv12_get_dst_fd() {
    return g_dst_fd;
}

// 获取 dst 的虚拟地址（供 Python 直接读取内存）
// 注意：返回的指针长度为 g_dst_size（初始化时分配的大小）
extern "C" void* nv12_get_dst_virt() {
    return g_dst_virt;
}

// 获取 dst buffer 大小（bytes）
extern "C" size_t nv12_get_dst_size() {
    return g_dst_size;
}

// 反初始化并释放全局 dst buffer
extern "C" void nv12_to_rgb_deinit() {
    if (g_dst_virt) {
        dma_buf_free((int)g_dst_size, &g_dst_fd, g_dst_virt);
        g_dst_virt = nullptr;
        g_dst_fd = -1;
        g_dst_size = 0;
    }
    if (g_rga_inited) {
        c_RkRgaDeInit();
        g_rga_inited = false;
    }
    g_max_w = g_max_h = 0;
}


